package learning.test;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.util.*;
import learning.crf.training.CRFParameters;
import learning.crf.training.CollinsTraining;
import learning.crf.training.Evaluator;
import learning.crf.training.EvaluatorCV;
import learning.crf.training.FormatHelper;
import learning.crf.training.Evaluator.Example;
import learning.data.Dataset;
import learning.data.SequenceDataReader;
import learning.data.document.SequenceDocument;
import learning.logging.ConsoleLogger;
import learning.logging.Log;
import learning.logging.Log.Level;

public class ZclCRF {
	List<String> trainfiles;
	List<String> testfiles;
	
	String unifiedTrainFile;
	String unifiedTestFile;
	
	static String infoboxClass = "comedian";
	static String attribute = "nationality";
	static String suffix = "";
	//static String suffix = "_rebalanced";
	
	static String file = "O:/unix/projects/pardosa/data10/raphaelh/learning/" + infoboxClass + "_" + attribute + suffix;
	//static String file = "O:/unix/projects/pardosa/data10/raphaelh/tmp2/a";

	///////TO WRITE/////////////////////
	static String extractorFile = "temp";
	
	public ZclCRF(List<String> trainfiles,List<String> testfiles)throws Exception{
		String unifiedTrainFile = merge(trainfiles);
		String unifiedTestFile = merge(testfiles);
		Dataset<SequenceDocument> dataset_train = SequenceDataReader.read(unifiedTrainFile);
		Dataset<SequenceDocument> dataset_test = SequenceDataReader.read(unifiedTestFile);
		train(dataset_train);
		test(dataset_test);
	}
	
	private String merge(List<String>trainfiles){
		return "";
	}
	
	public double calcF1(){
		return 0;
	}
	
	public static void main(String[] args) throws Exception {
		
		// load dataset
		System.out.println("reading data");
		Dataset<SequenceDocument> dataset = SequenceDataReader.read(file);
		
		train(dataset);
		test(dataset);
		cv(dataset);
		
		System.out.println("done training");
		
	}
	
	private static void train(Dataset<SequenceDocument> dataset) throws Exception {
		// train
		dataset.setIteratorView(10, 0, true);
		System.out.println("training");
		
		Log.setLogger(new ConsoleLogger());
		Log.setLevel(Level.FINE);
		
		CollinsTraining ct = new CollinsTraining();
		CRFParameters parameters = ct.train(dataset);
		
		FileOutputStream fos = new FileOutputStream(extractorFile);
		parameters.serialize(fos);
		fos.close();		
	}
	
	private static void test(Dataset<SequenceDocument> dataset) throws Exception {
		
		CRFParameters parameters = new CRFParameters();
		FileInputStream fis = new FileInputStream(extractorFile);
		parameters.deserialize(fis);
		fis.close();		
		
		// test
		dataset.setIteratorView(3, 0, false);
		CollinsTraining ct = new CollinsTraining();
		Evaluator evaluator = new Evaluator(learning.crf.training.Evaluator.Level.FINE);
		ct.test(dataset, parameters, evaluator);
		
		System.out.println("tokenErrorRate:    " + evaluator.getTokenErrors() + " / " + evaluator.getTotalTokens());
		System.out.println("documentErrorRate: " + evaluator.getDocumentErrorRate());
		
		//numDocs 3177
		//viewStart: 0
		//viewEnd: 318
		//tokenErrorRate:    27 / 7984
		//documentErrorRate: 0.03773585
	}
	
	//cross validation
	
	private static void cv(Dataset<SequenceDocument> dataset) throws Exception {
		
		CollinsTraining ct = new CollinsTraining();
		EvaluatorCV evaluator = new EvaluatorCV(learning.crf.training.Evaluator.Level.FINE);
		ct.crossValidation(dataset, 5, evaluator);

		
		System.out.println("tokenErrorRate:    " + evaluator.getTokenErrors() + " / " + evaluator.getTotalTokens());
		System.out.println("documentErrorRate: " + evaluator.getDocumentErrorRate());

		System.out.println("-------------------------------------- CORRECT:");
		for (Evaluator e : evaluator.getRunEvaluators()) {
			
			for (SequenceDocument doc : e.getCorrect()) {
				boolean containsExtraction = false;
				for (int i=0; i < doc.labels.length; i++)
					if (doc.labels[i] == 1) containsExtraction = true;
				
				if (containsExtraction)
					System.out.println(FormatHelper.formatDocument(doc));
			}
		}
		
		System.out.println("-------------------------------------- INCORRECT:");
		for (Evaluator e : evaluator.getRunEvaluators()) {
			
			for (Example ex : e.getErrors()) {
				System.out.println(FormatHelper.formatExample(ex.doc, ex.predictedParse, ex.trueParse));
			}
		}		
	}
}

